;; Copyright 2019 Google LLC
;;
;; Licensed under the Apache License, Version 2.0 (the "License");
;; you may not use this file except in compliance with the License.
;; You may obtain a copy of the License at
;;
;;     http://www.apache.org/licenses/LICENSE-2.0
;;
;; Unless required by applicable law or agreed to in writing, software
;; distributed under the License is distributed on an "AS IS" BASIS,
;; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
;; See the License for the specific language governing permissions and
;; limitations under the License.

(library (eval)
  (export eval read-and-eval)
  (import (rnrs)
          (rnrs mutable-pairs))

  (define (trace-value x)
    (write x)
    (newline)
    x)
  (define (eval-error x what)
    (trace-value x)
    (error 'eval what))

  (define (make-const c) `(const ,c))
  (define (const-value x) (cadr x))

  (define (make-ref idx) `(ref ,idx))
  (define (ref-idx x) (cadr x))

  (define (make-set idx val) `(set ,idx ,val))
  (define (set-idx x) (cadr x))
  (define (set-val x) (caddr x))

  (define (make-seq head tail) `(seq ,head ,tail))
  (define (seq-head x) (cadr x))
  (define (seq-tail x) (caddr x))

  (define (make-if test then else) `(if ,test ,then ,else))
  (define (if-test x) (cadr x))
  (define (if-then x) (caddr x))
  (define (if-else x) (cadddr x))

  (define (make-let vals body)
    `(let ,vals ,body))
  (define (let-vals x) (cadr x))
  (define (let-body x) (caddr x))

  (define (and-map p ls)
    (or (null? ls) (and (p (car ls)) (and-map p (cdr ls)))))

  (define (make-letrec vals body)
    (unless (and-map (lambda (x) (eq? (car x) 'lambda)) vals)
      (eval-error vals "bad letrec values"))
    `(letrec ,vals ,body))
  (define (letrec-vals x) (cadr x))
  (define (letrec-body x) (caddr x))

  (define (make-lambda nargs body)
    `(lambda ,nargs ,body))
  (define (lambda-nargs x) (cadr x))
  (define (lambda-body x) (caddr x))

  (define (make-call proc args)
    `(call ,proc . ,args))
  (define (call-proc x) (cadr x))
  (define (call-args x) (cddr x))

  (define (find-var env var idx)
    (when (null? env)
      (eval-error var "unbound variable"))
    (if (eq? var (car env))
        idx
        (find-var (cdr env) var (+ idx 1))))

  (define (lookup-var idx env)
    (if (eq? idx 0)
        (car env)
        (lookup-var (- idx 1) env)))

  (define (extend-env env val)
    (cons val env))
  (define (extend-env* env vals)
    (if (null? vals)
        env
        (extend-env* (extend-env env (car vals)) (cdr vals))))

  (define (expand expr env)
    (cond
     ((or (number? expr) (boolean? expr) (char? expr) (string? expr))
      (make-const expr))
     ((symbol? expr)
      (make-ref (find-var env expr 0)))
     ((pair? expr)
      (let ((head (car expr)))
        (cond
         ((eq? head 'quote)
          (make-const (cadr expr)))
         ((eq? head 'begin)
          (if (null? (cdr expr))
              (make-const (when #f #f))
              (expand-body (cdr expr) env)))
         ((eq? head 'if)
          (make-if (expand (cadr expr) env)
                   (expand (caddr expr) env)
                   (expand (cadddr expr) env)))
         ((eq? head 'let*)
          (let ((bindings (cadr expr))
                (body (cddr expr)))
            (if (null? bindings)
                body
                (expand `(let (,(car bindings))
                           (let* ,(cdr bindings) . ,body))
                        env))))
         ((eq? head 'let)
          (if (symbol? (cadr expr))
              (let ((name (cadr expr))
                    (vars (map car (caddr expr)))
                    (inits (map cadr (caddr expr)))
                    (body (cdddr expr)))
                ;; Named let.
                (expand `(letrec ((,name (lambda ,vars . ,body)))
                           (,name . ,inits))
                        env))
              (let ((vars (map car (cadr expr)))
                    (vals (map cadr (cadr expr)))
                    (body (cddr expr)))
                (unless (and-map symbol? vars)
                  (eval-error vars "bad let vars"))
                (make-let (expand-exprs vals env)
                          (expand-body body (extend-env* env vars))))))
         ((eq? head 'letrec)
          (let ((vars (map car (cadr expr)))
                (vals (map cadr (cadr expr)))
                (body (cddr expr)))
            (unless (and-map symbol? vars)
              (eval-error vars "bad letrec vars"))
            (let ((env (extend-env* env vars)))
              (make-let (map (lambda (var) (make-const #f)) vars)
                        (expand-letrec-inits vars (expand-exprs vals env)
                                             env
                                             (expand-body body env))))))
         ((eq? head 'lambda)
          (let ((vars (cadr expr))
                (body (cddr expr)))
            (unless (and-map symbol? vars)
              (eval-error vars "bad lambda vars"))
            (let ((nargs (length vars)))
              (make-lambda nargs
                           (expand-body body (extend-env* env vars))))))
         (else
          (make-call (expand head env) (expand-exprs (cdr expr) env))))))
     (else (eval-error expr "invalid expression"))))

  (define (expand-letrec-inits vars vals env tail)
    (if (null? vars)
        tail
        (let ((set (make-set (find-var env (car vars) 0) (car vals))))
          (expand-letrec-inits (cdr vars) (cdr vals) env (make-seq set tail)))))

  (define (expand-exprs exprs env)
    (map (lambda (expr) (expand expr env)) exprs))

  (define (expand-body exprs env)
    (when (null? exprs) (error 'eval "empty body"))
    (let ((head (expand (car exprs) env)))
      (if (null? (cdr exprs))
          head
          (make-seq head (expand-body (cdr exprs) env)))))

  (define (list-set! list n x)
    (set-car! (list-tail list n) x))

  (define (%eval x env)
    (let ((tag (car x)))
      (cond
       ((eq? tag 'const)
        (const-value x))
       ((eq? tag 'ref)
        (list-ref env (ref-idx x)))
       ((eq? tag 'set)
        (list-set! env (set-idx x) (%eval (set-val x) env)))
       ((eq? tag 'seq)
        (%eval (seq-head x) env)
        (%eval (seq-tail x) env))
       ((eq? tag 'if)
        (if (%eval (if-test x) env)
            (%eval (if-then x) env)
            (%eval (if-else x) env)))
       ((eq? tag 'let)
        (%eval (let-body x) (extend-env* env (%eval* (let-vals x) env))))
       ((eq? tag 'lambda)
        (let ((nargs (lambda-nargs x))
              (body (lambda-body x)))
          (lambda (args)
            (unless (eq? (length args) nargs)
              (eval-error args "wrong number of args to lambda"))
            (%eval body (extend-env* env args)))))
       ((eq? tag 'call)
        (let ((args (%eval* (call-args x) env)))
          ((%eval (call-proc x) env) args)))
       (else (eval-error x "unexpected expanded expression")))))

  (define (%eval* x* env)
    (if (null? x*)
        '()
        (let ((x (car x*)) (x* (cdr x*)))
          (cons (%eval x env)
                (%eval* x* env)))))

  (define (wrap/0 p)
    (lambda (args)
      (unless (eq? (length args) 0)
        (eval-error p "wrong number of arguments"))
      (p)))
  (define (wrap/1 p)
    (lambda (args)
      (unless (eq? (length args) 1)
        (eval-error p "wrong number of arguments"))
      (p (car args))))
  (define (wrap/2 p)
    (lambda (args)
      (unless (eq? (length args) 2)
        (eval-error p "wrong number of arguments"))
      (p (car args) (cadr args))))
  (define (wrap/3 p)
    (lambda (args)
      (unless (eq? (length args) 3)
        (eval-error p "wrong number of arguments"))
      (p (car args) (cadr args) (caddr args))))

  (define (default-environment)
    `((*                 . ,(wrap/2 *))
      (+                 . ,(wrap/2 +))
      (-                 . ,(wrap/2 -))
      (<                 . ,(wrap/2 <))
      (>                 . ,(wrap/2 >))
      (append            . ,(wrap/2 append))
      (assq              . ,(wrap/2 assq))
      (bitwise-and       . ,(wrap/2 bitwise-and))
      (bitwise-arithmetic-shift-left  . ,(wrap/2 bitwise-arithmetic-shift-left))
      (bitwise-arithmetic-shift-right . ,(wrap/2 bitwise-arithmetic-shift-right))
      (bitwise-ior       . ,(wrap/2 bitwise-ior))
      (bitwise-not       . ,(wrap/1 bitwise-not))
      (boolean?          . ,(wrap/1 boolean?))
      (caaar             . ,(wrap/1 caaar))
      (caadar            . ,(wrap/1 caadar))
      (caaddr            . ,(wrap/1 caaddr))
      (caadr             . ,(wrap/1 caadr))
      (caar              . ,(wrap/1 caar))
      (cadadr            . ,(wrap/1 cadadr))
      (cadar             . ,(wrap/1 cadar))
      (caddar            . ,(wrap/1 caddar))
      (cadddr            . ,(wrap/1 cadddr))
      (caddr             . ,(wrap/1 caddr))
      (cadr              . ,(wrap/1 cadr))
      (car               . ,(wrap/1 car))
      (cdaddr            . ,(wrap/1 cdaddr))
      (cdadr             . ,(wrap/1 cdadr))
      (cdar              . ,(wrap/1 cdar))
      (cddar             . ,(wrap/1 cddar))
      (cdddr             . ,(wrap/1 cdddr))
      (cddr              . ,(wrap/1 cddr))
      (cdr               . ,(wrap/1 cdr))
      (char->integer     . ,(wrap/1 char->integer))
      (char-ci<?         . ,(wrap/2 char-ci<?))
      (char-numeric?     . ,(wrap/1 char-numeric?))
      (char-whitespace?  . ,(wrap/1 char-whitespace?))
      (char?             . ,(wrap/1 char?))
      (cons              . ,(wrap/2 cons))
      (display           . ,(wrap/1 display))
      (div0              . ,(wrap/2 div0))
      (eof-object        . ,(wrap/0 eof-object))
      (eq?               . ,(wrap/2 eq?))
      (error             . ,(wrap/2 error))
      (fold-left         . ,(wrap/3 fold-left))
      (fold-right        . ,(wrap/3 fold-right))
      (integer->char     . ,(wrap/1 integer->char))
      (length            . ,(wrap/1 length))
      (list->string      . ,(wrap/1 list->string))
      (list-ref          . ,(wrap/2 list-ref))
      (list-tail         . ,(wrap/2 list-tail))
      (map               . ,(wrap/2 map))
      (max               . ,(wrap/2 max))
      (mod0              . ,(wrap/2 mod0))
      (newline           . ,(wrap/0 newline))
      (null?             . ,(wrap/1 null?))
      (number?           . ,(wrap/1 number?))
      (pair?             . ,(wrap/1 pair?))
      (peek-char         . ,(wrap/0 peek-char))
      (procedure?        . ,(wrap/1 procedure?))
      (read              . ,(wrap/0 read))
      (read-char         . ,(wrap/0 read-char))
      (set-car!          . ,(wrap/2 set-car!))
      (set-cdr!          . ,(wrap/2 set-cdr!))
      (string->list      . ,(wrap/1 string->list))
      (string->symbol    . ,(wrap/1 string->symbol))
      (string?           . ,(wrap/1 string?))
      (symbol->string    . ,(wrap/1 symbol->string))
      (symbol?           . ,(wrap/1 symbol?))
      (write             . ,(wrap/1 write))
      (write-char        . ,(wrap/1 write-char))
      (zero?             . ,(wrap/1 zero?))))

  (define (eval expr)
    (let ((env (default-environment)))
      (%eval (expand expr (map car env)) (map cdr env))))
  (define (read-and-eval)
    (eval (read))))
